package com.luckzj.sec.bridge;

import com.luckzj.sec.config.BridgeConfig;
import io.netty.bootstrap.Bootstrap;
import io.netty.bootstrap.ServerBootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioServerSocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.Executors;

/**
 * Abstract Bridge
 * Implement basic netty operations
 * subclasses should use coder and decoders to implement http/ftp bridge functions
 */
public abstract class AbstractBridge {
    private static Logger logger = LoggerFactory.getLogger(AbstractBridge.class);

    private BridgeConfig bridgeConfig;
    private BridgeMapper bridgeMapper = new BridgeMapper();

    public AbstractBridge(BridgeConfig bridgeConfig) {
        this.bridgeConfig = bridgeConfig;
    }

    /**
     * Start a tcp server and transfer network flow to destination
     */
    public void startBridge() {
        Executors.newSingleThreadExecutor().execute(
                new Runnable() {
                    @Override
                    public void run() {
                        startServer();
                    }
                }
        );
    }

    private void startServer() {
        EventLoopGroup bossGrp = new NioEventLoopGroup();
        EventLoopGroup workerGrp = new NioEventLoopGroup();

        ServerBootstrap bootstrap = new ServerBootstrap();
        bootstrap.group(bossGrp, workerGrp)
                .channel(NioServerSocketChannel.class)
                .option(ChannelOption.SO_BACKLOG, 128)
                .childOption(ChannelOption.SO_KEEPALIVE, true)
                .childHandler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) throws Exception {
                        onSetupServerPipeline(ch);
                    }
                });

        try {
            ChannelFuture channelFuture = bootstrap.bind(bridgeConfig.getSrcPort()).sync();

            if (logger.isInfoEnabled()) {
                logger.info("server started for port {}", bridgeConfig.getSrcPort());
            }

            channelFuture.channel().closeFuture().sync();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } finally {
            bossGrp.shutdownGracefully();
            workerGrp.shutdownGracefully();
        }
    }

    protected abstract void setupServerChannel(SocketChannel ch);

    public void onSetupServerPipeline(SocketChannel ch) {
        setupServerChannel(ch);

        ch.pipeline().addLast(new ChannelInboundHandlerAdapter() {
            @Override
            public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
                onServerChannelRead(ctx, msg);

                // super.channelRead(ctx, msg);
            }

            @Override
            public void channelActive(ChannelHandlerContext ctx) throws Exception {
                onServerChannelActive(ctx);

                super.channelActive(ctx);
            }

            @Override
            public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
                onServerExceptionCaught(ctx, cause);

                super.exceptionCaught(ctx, cause);
            }
        });
    }

    public void onServerChannelActive(ChannelHandlerContext ctx) {
        if (logger.isInfoEnabled()) {
            logger.info("server active from {}" + ctx.channel().remoteAddress());
        }

        if (bridgeMapper.getDst(ctx.channel()) != null) {
            if (logger.isWarnEnabled()) {
                logger.warn("server channel already mapped. ch = {}", ctx.channel().localAddress());
            }

            bridgeMapper.removeSrc(ctx.channel());
        }

        // create new client channel
        Channel client = startClient(bridgeConfig.getDstHost(), bridgeConfig.getDstPort());
        bridgeMapper.addMap(ctx.channel(), client);

        if (logger.isInfoEnabled()) {
            logger.info("create new connection, src port = {}, dst = {}:{}", bridgeConfig.getSrcPort()
                    , bridgeConfig.getDstHost(), bridgeConfig.getDstPort());
        }
    }

    public void onServerChannelRead(ChannelHandlerContext ctx, Object msg) {
        logger.info("server channel read.");
        Channel client = bridgeMapper.getDst(ctx.channel());
        if (client == null) {
            if (logger.isErrorEnabled()) {
                logger.error("can not find a dst channel! srcPort = {}", bridgeConfig.getSrcPort());
            }

            return;
        }

        // proxy this msg
        client.writeAndFlush(msg);
    }

    public void onServerExceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        // close server and client connection
        Channel dst = bridgeMapper.getDst(ctx.channel());
        bridgeMapper.removeSrc(ctx.channel());

        try {
            if (dst != null) {
                dst.closeFuture().sync();
            }
            ctx.channel().closeFuture().sync();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    private Channel startClient(String host, int port) {
        Bootstrap b = new Bootstrap();
        b.group(new NioEventLoopGroup());
        b.channel(NioSocketChannel.class); // (3)
        b.option(ChannelOption.SO_KEEPALIVE, true); // (4)
        b.handler(new ChannelInitializer<SocketChannel>() {
            @Override
            public void initChannel(SocketChannel ch) throws Exception {
                onSetupClientPipeline(ch);
            }
        });

        // Start the client.
        ChannelFuture f = null; // (5)
        try {
            f = b.connect(host, port).sync();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        return f.channel();
    }

    protected abstract void setupClientChannel(SocketChannel ch);

    public void onSetupClientPipeline(SocketChannel ch) {
        setupClientChannel(ch);

        ch.pipeline().addLast(new ChannelInboundHandlerAdapter() {
            @Override
            public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
                onClientChannelRead(ctx, msg);
                // super.channelRead(ctx, msg);
            }

            @Override
            public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
                onClientExceptionCaught(ctx, cause);

                super.exceptionCaught(ctx, cause);
            }
        });
    }

    private void onClientChannelRead(ChannelHandlerContext ctx, Object msg) {
        logger.info("client channel read.");
        Channel server = bridgeMapper.getSrc(ctx.channel());
        if (server == null) {
            if (logger.isErrorEnabled()) {
                logger.error("can not find a dst channel! srcPort = {}", bridgeConfig.getSrcPort());
            }

            return;
        }

        // proxy this msg
        server.writeAndFlush(msg);
    }

    public void onClientExceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        // close server and client connection
        Channel src = bridgeMapper.getSrc(ctx.channel());
        bridgeMapper.removeDst(ctx.channel());

        try {
            src.closeFuture().sync();
            ctx.channel().closeFuture().sync();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    public BridgeConfig getBridgeConfig() {
        return bridgeConfig;
    }
}
